#include <vector>
#include <memory>
#include "dense_layer.h"

class BranchNetwork {
public:
    std::vector<DenseLayer> layers;
    std::shared_ptr<MatrixOperations> matrixOperations;
    std::shared_ptr<ActivationFunction> activationFunction;

    /**
     * 构造函数初始化一个分支网络 (BranchNetwork)
     * 
     * @param layerSizes 各层的大小，每个元素代表一个层的神经元数目。
     * @param lr 学习率，用于在训练过程中更新权重。
     * @param matOp 矩阵运算的共享指针，用于执行所有矩阵相关的操作，比如前向传播和反向传播。
     * @param actFunc 激活函数的共享指针，用于每个全连接层。
     */
    BranchNetwork(const std::vector<size_t>& layerSizes, float lr, std::shared_ptr<MatrixOperations> matOp, std::shared_ptr<ActivationFunction> actFunc)
        : matrixOperations(matOp), activationFunction(actFunc) {
        // 根据指定的层数和各层大小创建全连接层
        if (layerSizes.size() < 2) {
            throw std::invalid_argument("layerSizes must contain at least two elements for input and output sizes.");
        }
        
        size_t prevLayerSize = layerSizes.front();
        for (size_t i = 1; i < layerSizes.size(); ++i) {
            layers.emplace_back(DenseLayer(prevLayerSize, layerSizes[i], lr, activationFunction, matrixOperations));
            prevLayerSize = layerSizes[i]; // 更新前一层的大小为当前层的大小
        }
    }

    std::vector<std::vector<float>> forward(const std::vector<std::vector<float>>& inputs) {
        // 前向传播，从输入层到输出层
        std::vector<std::vector<float>> layerOutputs = inputs;
        for (auto& layer : layers) {
            layerOutputs = layer.forward(layerOutputs);
        }
        return layerOutputs;
    }

    void backward(const std::vector<std::vector<float>>& outputGradients) {
        // 反向传播，从输出层到输入层
        std::vector<std::vector<float>> gradients = outputGradients;
        for (auto it = layers.rbegin(); it != layers.rend(); ++it) {
            gradients = it->backward(gradients);
        }
    }
};
